import torch
import numpy as np

def calculate_apoz(activations):
    """计算APoZ值（3.2.1节）"""
    return (activations == 0).float().mean(dim=0)

def personalized_pruning(model, train_data, threshold=0.5):
    """
    个性化剪枝（3.2.1节）
    基于APoZ值的神经元剪枝
    """
    model.eval()
    activations = {}
    
    # 注册钩子收集激活值
    hooks = []
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.ReLU):
            hook = module.register_forward_hook(
                lambda m, i, o: activations.setdefault(name, []).append(o)
            )
            hooks.append(hook)
    
    # 前向传播收集激活
    with torch.no_grad():
        for x, _ in train_data:
            model(x)
    
    # 移除钩子
    for hook in hooks:
        hook.remove()
    
    # 计算APoZ并剪枝
    for name, acts in activations.items():
        all_acts = torch.cat(acts, dim=0)
        apoz = calculate_apoz(all_acts)
        
        # 剪枝低激活神经元
        for i, score in enumerate(apoz):
            if score > threshold:
                # 实际实现中需要修改权重矩阵
                print(f"Pruning neuron {i} in layer {name} with APoZ {score:.4f}")

def dynamic_pruning(models, metrics, threshold=0.05):
    """
    Algorithm 1: Dynamic Pruning Algorithm
    基于模型稳定性的客户端选择
    """
    selected_clients = []
    
    for client_id, model in enumerate(models):
        if client_id not in metrics:
            continue
            
        acc_history = metrics[client_id]['accuracy']
        if len(acc_history) < 2:
            continue
            
        delta_acc = abs(acc_history[-1] - acc_history[-2])
        if delta_acc < threshold:
            selected_clients.append(client_id)
    
    print(f"Selected {len(selected_clients)}/{len(models)} clients")
    return selected_clients
